from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import Response
from PIL import Image
import io
import logging
from processing import remove_background

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="人像背景移除服务",
    description="将人像照片背景转为白底的API服务",
    version="1.0.0"
)

# pip install python-multipart

@app.post("/remove-background", summary="移除背景并转为白底")
async def remove_image_background(file: UploadFile = File(..., description="上传的人像照片")):
    # 验证文件类型
    allowed_extensions = ['.png', '.jpg', '.jpeg', '.webp']
    file_extension = file.filename[file.filename.rfind('.'):].lower()

    if file_extension not in allowed_extensions:
        logger.error(f"不支持的文件格式: {file.filename}")
        raise HTTPException(
            status_code=400,
            detail=f"不支持的文件格式，支持的格式: {', '.join(allowed_extensions)}"
        )

    try:
        logger.info(f"开始处理图片: {file.filename}")

        # 读取上传的图片
        contents = await file.read()
        input_image = Image.open(io.BytesIO(contents))

        # 处理图片
        output_image = remove_background(input_image)

        # 将结果转为字节流
        img_byte_arr = io.BytesIO()
        output_image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        logger.info(f"图片处理完成: {file.filename}")
        return Response(content=img_byte_arr, media_type="image/png")

    except Exception as e:
        logger.error(f"图片处理失败: {file.filename}, 错误: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"图片处理失败: {str(e)}"
        )


@app.get("/health", summary="服务健康检查")
def health_check():
    return {"status": "healthy", "version": "1.0.0"}